package com.chukun.flink.stream.state.key;

import org.apache.flink.api.common.functions.AggregateFunction;
import org.apache.flink.api.common.functions.RichFlatMapFunction;
import org.apache.flink.api.common.state.AggregatingState;
import org.apache.flink.api.common.state.AggregatingStateDescriptor;
import org.apache.flink.api.common.typeinfo.TypeHint;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.KeyedStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.util.Collector;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * @author chukun
 * @version 1.0.0
 * @description 带状态的flatmap操作符 AggregatingState的基本使用
 * @createTime 2022年05月15日 22:59:00
 */
public class AggregatingStateFlatMap extends RichFlatMapFunction<Tuple2<Integer, Integer>, Tuple2<Integer,Integer>> {

    private static final Logger logger = LoggerFactory.getLogger(AggregatingStateFlatMap.class);

    // 创建AggregatingState类型的状态结构
    private transient AggregatingState<Tuple2<Integer,Integer>,Double> aggregatingState;

    @Override
    public void flatMap(Tuple2<Integer, Integer> input, Collector<Tuple2<Integer, Integer>> collector) throws Exception {
        aggregatingState.add(input);
        Double value = aggregatingState.get();
        if (value > 10) {
            // 输出该值
            logger.info("current value : {}", value);
            // 重置状态值
            aggregatingState.clear();
            collector.collect(Tuple2.of(value.intValue(), value.intValue()));
        }
    }

    @Override
    public void open(Configuration parameters) throws Exception {
        logger.info("{},{}", Thread.currentThread().getName(), "恢复或初始化状态");

        // 创建一个AggregatingStateDescriptor状态描述符，指定名称为AggregatingStateFlatMap，类型为 Tuple2<Integer, Integer>
        AggregatingStateDescriptor descriptor = new AggregatingStateDescriptor(
                "AggregatingStateFlatMap", new AggregateFunction<Tuple2<Integer,Integer>, AverageAccumulator, Double>() {
            @Override
            public AverageAccumulator createAccumulator() {
                return new AverageAccumulator();
            }

            @Override
            public AverageAccumulator add(Tuple2<Integer, Integer> value, AverageAccumulator accumulator) {
                accumulator.add(value.f1);
                return accumulator;
            }

            @Override
            public Double getResult(AverageAccumulator accumulator) {
                return accumulator.getLocalValue();
            }

            @Override
            public AverageAccumulator merge(AverageAccumulator a, AverageAccumulator b) {
                a.add(b.getCount(), b.getSum());
                return a;
            }
        },
                TypeInformation.of(new TypeHint<Tuple2<Integer, Integer>>() {
                }));
        // 初始化valueState
        aggregatingState = getRuntimeContext().getAggregatingState(descriptor);
        logger.info("初始化 valueState ： {}", aggregatingState);
    }

    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        // 封装了数据源函数，以及keyBy操作的数据流
        KeyedStream<Tuple2<Integer, Integer>, Integer> keyedStream = KeyStateBase.before(env);

        // 在keyedStream中使用有状态的FlatMap操作符
        DataStream<Tuple2<Integer, Integer>> resultStream = keyedStream.flatMap(new AggregatingStateFlatMap());

        resultStream.print("输出结果 : ");

        env.execute("AggregatingStateFlatMap");
    }
}
